package io.sqrtqiezi.practise.spark.sort.secondary;

import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.PairFunction;
import scala.Tuple2;

import org.apache.spark.api.java.JavaSparkContext;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class SecondarySort {
    public static void main(String[] args) {
        if (args.length < 1) {
            System.err.println("Usage: SecondarySort <file>");
            System.exit(1);
        }

        final String inputPath = args[0];
        System.out.println("args[0]: <file>=" + inputPath);

        final JavaSparkContext ctx = new JavaSparkContext("local[*]", "secondary-sort");
        ctx.setLogLevel("WARN");

        final JavaRDD<String> lines = ctx.textFile(inputPath, 1);

        // 从 JavaRDD 创建键值对
        final JavaPairRDD<String, Tuple2<Integer, Integer>> pairs =
                lines.mapToPair(new PairFunction<String, String, Tuple2<Integer, Integer>>() {
                    @Override
                    public Tuple2<String, Tuple2<Integer, Integer>> call(String s) throws Exception {
                        final String[] tokens = s.split(",");
//                        System.out.println(tokens[0] + "," + tokens[1] + "," + tokens[2]);
                        final Integer time = new Integer(tokens[1]);
                        final Integer value = new Integer(tokens[2]);
                        final Tuple2<Integer, Integer> timeValue = new Tuple2<>(time, value);
                        return new Tuple2<>(tokens[0], timeValue);
                    }
                });

        // 验证上一步
        final List<Tuple2<String, Tuple2<Integer, Integer>>> output = pairs.collect();
        for (Tuple2<String, Tuple2<Integer, Integer>> item : output) {
            final String name = item._1;
            final Tuple2<Integer, Integer> timeValue = item._2;
            System.out.println(name + "," + timeValue._1 + "," + timeValue._2);
        }

        // 按 name 进行分组
        final JavaPairRDD<String, Iterable<Tuple2<Integer, Integer>>> groups = pairs.groupByKey();

        System.out.println("======= DEBUG1 =======");
        printGrouped(groups);

        // 在内存中对 reducer 值排序
        final JavaPairRDD<String, Iterable<Tuple2<Integer, Integer>>> sorted =
                groups.mapValues(new Function<
                        Iterable<Tuple2<Integer, Integer>>,
                        Iterable<Tuple2<Integer, Integer>>
                        >() {
                    @Override
                    public Iterable<Tuple2<Integer, Integer>> call(Iterable<Tuple2<Integer, Integer>> v1) throws Exception {
                        List<Tuple2<Integer, Integer>> newList = new ArrayList<>(iterableToList(v1));
                        Collections.sort(newList, new TupleComparator());
                        return newList;
                    }
                });

        // 打印最终结果
        System.out.println("======= DEBUG2 =======");
        printGrouped(sorted);

        ctx.close();
        System.exit(0);
    }

    private static void printGrouped(JavaPairRDD<String, Iterable<Tuple2<Integer, Integer>>> grouped) {
        final List<Tuple2<String, Iterable<Tuple2<Integer, Integer>>>> output = grouped.collect();
        for (Tuple2<String, Iterable<Tuple2<Integer, Integer>>> item : output) {
            final Iterable<Tuple2<Integer, Integer>> list = item._2;
            System.out.println(item._1);
            for (Tuple2<Integer, Integer> item2 : list) {
                System.out.println(item2._1 + "," + item2._2);
            }
            System.out.println("=====================");
        }
    }

    static List<Tuple2<Integer,Integer>> iterableToList(Iterable<Tuple2<Integer,Integer>> iterable) {
        List<Tuple2<Integer,Integer>> list = new ArrayList<>();
        for (Tuple2<Integer,Integer> item : iterable) {
            list.add(item);
        }
        return list;
    }
}
